#!/bin/bash
# PostgreSQL从库设置脚本 - RedFire Stage2
# 自动配置PostgreSQL流复制

set -e

echo "=== PostgreSQL从库设置开始 ==="

# 环境变量检查
: ${POSTGRES_PRIMARY_HOST:="postgresql-primary"}
: ${POSTGRES_PRIMARY_PORT:="5432"}
: ${POSTGRES_REPLICA_USER:="replicator"}
: ${POSTGRES_REPLICA_PASSWORD:=""}
: ${POSTGRES_DB:="redfire_stage2"}
: ${POSTGRES_USER:="redfire_prod"}
: ${POSTGRES_PASSWORD:=""}

if [ -z "$POSTGRES_REPLICA_PASSWORD" ] || [ -z "$POSTGRES_PASSWORD" ]; then
    echo "错误: 复制用户密码或数据库密码未设置"
    exit 1
fi

# 等待主库启动
echo "等待主库 $POSTGRES_PRIMARY_HOST:$POSTGRES_PRIMARY_PORT 启动..."
while ! pg_isready -h $POSTGRES_PRIMARY_HOST -p $POSTGRES_PRIMARY_PORT -U $POSTGRES_USER; do
    echo "主库未就绪，等待5秒..."
    sleep 5
done

echo "主库已就绪，开始配置从库..."

# 停止PostgreSQL服务
pg_ctl stop -D $PGDATA -m fast || true

# 清理数据目录
rm -rf $PGDATA/*

# 创建复制用户 (在主库上执行)
echo "在主库上创建复制用户..."
PGPASSWORD=$POSTGRES_PASSWORD psql -h $POSTGRES_PRIMARY_HOST -p $POSTGRES_PRIMARY_PORT -U $POSTGRES_USER -d $POSTGRES_DB -c "
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '$POSTGRES_REPLICA_USER') THEN
        CREATE USER $POSTGRES_REPLICA_USER WITH REPLICATION ENCRYPTED PASSWORD '$POSTGRES_REPLICA_PASSWORD';
    END IF;
END
\$\$;
"

# 使用pg_basebackup创建从库
echo "使用pg_basebackup创建从库..."
PGPASSWORD=$POSTGRES_REPLICA_PASSWORD pg_basebackup \
    -h $POSTGRES_PRIMARY_HOST \
    -p $POSTGRES_PRIMARY_PORT \
    -U $POSTGRES_REPLICA_USER \
    -D $PGDATA \
    -Fp -Xs -P -R

# 创建standby.signal文件
touch $PGDATA/standby.signal

# 配置recovery.conf (PostgreSQL 12+使用postgresql.conf)
echo "配置从库连接参数..."
cat >> $PGDATA/postgresql.conf << EOF

# 从库复制配置
primary_conninfo = 'host=$POSTGRES_PRIMARY_HOST port=$POSTGRES_PRIMARY_PORT user=$POSTGRES_REPLICA_USER password=$POSTGRES_REPLICA_PASSWORD application_name=standby1'
primary_slot_name = 'standby_slot_1'
restore_command = 'cp /var/lib/postgresql/archive/%f %p'
archive_cleanup_command = 'pg_archivecleanup /var/lib/postgresql/archive %r'
EOF

# 在主库上创建复制槽
echo "在主库上创建复制槽..."
PGPASSWORD=$POSTGRES_PASSWORD psql -h $POSTGRES_PRIMARY_HOST -p $POSTGRES_PRIMARY_PORT -U $POSTGRES_USER -d $POSTGRES_DB -c "
SELECT pg_create_physical_replication_slot('standby_slot_1');
" || echo "复制槽可能已存在，继续..."

# 设置正确的权限
chown -R postgres:postgres $PGDATA
chmod 700 $PGDATA

echo "=== PostgreSQL从库设置完成 ==="
echo "从库将在容器启动时自动连接到主库"
echo "主库: $POSTGRES_PRIMARY_HOST:$POSTGRES_PRIMARY_PORT"
echo "复制用户: $POSTGRES_REPLICA_USER"
echo "复制槽: standby_slot_1"
